from functools import partial
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy.dialects.mysql import TIMESTAMP, DATETIME
from sqlalchemy.orm import sessionmaker, scoped_session, class_mapper
from config import MYSQL_DATABASE, MYSQL_HOST, MYSQL_PASSWORD,\
    MYSQL_USER, MYSQL_PORT, MYSQL_DATABASE


NotNullColumn = partial(Column, nullable=False, server_default='')


def model2dict(model):
    if not model:
        return {}
    fields = class_mapper(model.__class__).columns.keys()
    return dict((col, getattr(model, col)) for col in fields)


def model_to_dict(func):
    def wrap(*args, **kwargs):
        ret = func(*args, **kwargs)
        return model2dict(ret)
    return wrap


def models_to_list(func):
    def wrap(*args, **kwargs):
        ret = func(*args, **kwargs)
        return [model2dict(r) for r in ret]
    return wrap


class declare_base(object):

    # 设置通用字段
    create_time = NotNullColumn(DATETIME)
    update_time = NotNullColumn(TIMESTAMP, server_default='CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP')

    @declared_attr
    def __table_args__(cls):  # 设置数据库默认信息
        return {
            'mysql_engine': 'InnoDB',
            'mysql_charset': 'utf8'
        }


Base = declarative_base(cls=declare_base)  # 继承自己编写的通用cls


CONF = {
    'kwargs': {
        'pool_recycle': 3600,
        'pool_size': 5
    },
    'schema': 'mysql+pymysql://%s:%s@%s:%d/%s?charset=%s'
}

MYSQL_CONFIG = {
        "user": MYSQL_USER,
        "password": MYSQL_PASSWORD,
        "host": MYSQL_HOST,
        "port": MYSQL_PORT,
        "database":MYSQL_DATABASE,
        "charset":"utf8mb4"
    }

def create_session(user, password, host, port, database, charset):
    schema = CONF['schema'] % (user, password, host, port, database, charset)
    engine = create_engine(schema, **CONF['kwargs'])
    session = scoped_session(sessionmaker(bind=engine))
    return session()


class SQLStore(object):

    def __init__(self, db):
        self.master = create_session(**db)

    def execute(self, clause, params=None):
        try:
            return self.master.execute(clause, params=params)
        except Exception:
            self.rollback()
            raise

    def add(self, one):
        self.master.add(one)
        self.master.commit()
        return one

    def commit(self):
        return self.master.commit()

    def rollback(self):
        return self.master.rollback()

    def is_testing(self):
        return False


db = SQLStore(MYSQL_CONFIG)
